# 导入FastMCP框架
import json
from fastmcp import FastMCP
from pymongo import MongoClient
from config import MONGO_URI, MONGO_DB_NAME

# Connect to MongoDB
mongo = MongoClient(MONGO_URI)
db = mongo[MONGO_DB_NAME]

# 初始化FastMCP实例
mcp = FastMCP("TireMcp")


@mcp.tool()
async def search_products(keyword: str, brand: str, min_stock: float):
    """
    Get product list by searching for product name / code keywords, brand keywords or minimum inventory quantity

    Args:
        keyword: Product name/Code keyword
        brand: Brand name keyword
        min_stock: Minimum stock
    """
    filters = {}

    if keyword:
        filters["$or"] = [{"name": {"$regex": keyword, "$options": "i"}}, {"code": {"$regex": keyword, "$options": "i"}}]

    if brand:
        filters["brand_info.name"] = {"$regex": brand, "$options": "i"}

    if min_stock and min_stock > 0:
        filters["stock_summary.total_stock"] = {"$gt": min_stock}

    results = db.products.find(
        filters, {"_id": 0, "name": 1, "code": 1, "stock_summary.total_stock": 1, "brand_info.name": 1, "images": 1}
    ).limit(10)

    data = []
    for r in results:
        images = r.get("images") or []
        image_url = images[0].get("url") if images and isinstance(images[0], dict) else ""
        data.append(
            {
                "name": r.get("name", ""),
                "code": r.get("code", ""),
                "stock": r.get("stock_summary", {}).get("total_stock", 0),
                "brand": r.get("brand_info", {}).get("name", ""),
                "image": image_url,
            }
        )

    return json.dumps(data, ensure_ascii=False)


@mcp.tool()
async def get_all_brands():
    """
    Get all unique brand names from products
    获取所有产品品牌
    """
    filters = {}
    results = db.products.find(filters, {"_id": 0, "brand_info.name": 1})

    # 使用集合自动去重
    unique_names = {doc["brand_info"]["name"] for doc in results if doc.get("brand_info")}

    data = [{"name": name} for name in unique_names]
    return json.dumps(data, ensure_ascii=False)


@mcp.tool()
async def get_details_by_product(brand: str) -> str:
    """
    获取(这个品牌下)最受欢迎的几款产品

    参数:
        brand (str): 品牌名称

    返回:
        str: 这个品牌下最受欢迎的几款产品
    """
    three_brand_name = []
    if brand:
        three_brand_name = [brand]
    else:
        results = db.products.find({}, {"_id": 0, "brand_info.name": 1})
        # 使用集合自动去重
        unique_names = {doc["brand_info"]["name"] for doc in results if doc.get("brand_info")}
        brand_name = [name for name in unique_names]
        three_brand_name = brand_name[:3]  # 随机3个品牌

    data = []
    for brand in three_brand_name:
        filters = {"brand_info.name": {"$regex": brand, "$options": "i"}}

        results = db.products.find(
            filters, {"_id": 0, "name": 1, "code": 1, "stock_summary.total_stock": 1, "brand_info.name": 1, "images": 1}
        ).limit(3)

        for r in results:
            images = r.get("images") or []
            image_url = images[0].get("url") if images and isinstance(images[0], dict) else ""
            data.append(
                {
                    "name": r.get("name", ""),
                    "code": r.get("code", ""),
                    "stock": r.get("stock_summary", {}).get("total_stock", 0),
                    "brand": r.get("brand_info", {}).get("name", ""),
                    "image": image_url,
                }
            )

    return json.dumps(data, ensure_ascii=False)


# 程序入口点
if __name__ == "__main__":
    # 启动MCP服务器
    mcp.settings.host = "0.0.0.0"
    mcp.settings.port = 9001
    mcp.run(transport="sse")
